package com.c2b.algorithm.leetcode.base;

import java.util.ArrayList;
import java.util.List;
import java.util.PriorityQueue;
import java.util.Random;

/**
 * <a href="https://leetcode.cn/problems/kth-largest-element-in-an-array/description/?envType=study-plan-v2&envId=top-interview-150">数组中的第K个最大元素(Kth Largest Element in an Array)</a>
 * <p>给定整数数组 nums 和整数 k，请返回数组中第 k 个最大的元素。</p>
 * <p>请注意，你需要找的是数组排序后的第 k 个最大的元素，而不是第 k 个不同的元素。</p>
 * <p>你必须设计并实现时间复杂度为 O(n) 的算法解决此问题。</p>
 *
 * <p>
 * <b>示例：</b>
 * <pre>
 * 示例 1:
 *      输入: [3,2,1,5,6,4], k = 2
 *      输出: 5
 *
 * 示例 2:
 *      输入: [3,2,3,1,2,4,5,5,6], k = 4
 *      输出: 4
 * </pre>
 * </p>
 *
 * <p>
 * <b>提示：</b>
 *  <ul>
 *      <li>1 <= k <= nums.length <= 10^5</li>
 *      <li>-10^4 <= nums[i] <= 10^4</li>
 *  </ul>
 * </p>
 *
 * @author c2b
 * @since 2023/10/20 16:13
 */
public class LC0215KthLargestElementInArray_M {

    static class Solution {

        /**
         * 快速排序。
         * <p>
         * 快速排序的核心包括“哨兵划分” 和 “递归” 。
         *     <ul>
         *         <li>哨兵划分： 以数组某个元素（一般选取首元素）为基准数，将所有小于基准数的元素移动至其左边，大于基准数的元素移动至其右边。</li>
         *         <li>递归： 对 左子数组 和 右子数组 递归执行 哨兵划分，直至子数组长度为 1 时终止递归，即可完成对整个数组的排序。</li>
         *     </ul>
         * </p>
         */
        public int findKthLargest(int[] nums, int k) {
            List<Integer> numList = new ArrayList<>();
            for (int num : nums) {
                numList.add(num);
            }
            return quickSelect(numList, k);
        }

        private int quickSelect(List<Integer> numList, int k) {
            // 随机选择基准数
            Random random = new Random();
            int pivot = numList.get(random.nextInt(numList.size()));
            // 将大于、小于、等于 pivot 的元素 划分至 不同的集合中
            List<Integer> smallerList = new ArrayList<>();
            List<Integer> biggerList = new ArrayList<>();
            for (int num : numList) {
                if (num > pivot) {
                    biggerList.add(num);
                } else if (num < pivot) {
                    smallerList.add(num);
                }
            }
            // 第 k 大元素在 biggerList 中，递归划分
            if (biggerList.size() >= k) {
                return quickSelect(biggerList, k);
            }
            // 第 k 大元素在 smallerList 中，递归划分
            if (numList.size() - smallerList.size() < k) {
                return quickSelect(smallerList, k + smallerList.size() - numList.size());
            }
            // 第 k 大元素在 equalList 中，直接返回 pivot
            return pivot;
        }

        public int findKthLargest2(int[] nums, int k) {
            // 维护一个小根堆
            PriorityQueue<Integer> queue = new PriorityQueue<>(k);
            for (int num : nums) {
                if (queue.size() < k) {
                    queue.offer(num);
                } else {
                    if (num > queue.peek()) {
                        queue.poll();
                        queue.offer(num);
                    }
                }
            }
            return queue.peek();
        }
    }

    public static void main(String[] args) {
        Solution solution = new Solution();
        System.out.println(solution.findKthLargest(new int[]{3, 2, 1, 5, 6, 4}, 2));
        System.out.println(solution.findKthLargest(new int[]{3, 2, 3, 1, 2, 4, 5, 5, 6}, 4));
    }
}
